from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage

def retain_messages(
    messages: "list[ModelMessage]",
) -> "list[ModelMessage]":
    return messages

chat_client = Agent(
    "openai:o4-mini",
    system_prompt="You are a helpful assistant.",
    history_processors=[retain_messages],
)

import os

from dotenv import load_dotenv
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage
from shiny.express import ui

_ = load_dotenv()


def retain_messages(
    messages: "list[ModelMessage]",
) -> "list[ModelMessage]":
    return messages


chat_client = Agent(
    "openai:gpt-4.1-nano-2025-04-14",
    system_prompt="You are a helpful assistant.",
    history_processors=[retain_messages],
)


# Set some Shiny page options
ui.page_opts(
    title="OpenAI with Pydantic AI",
    fillable=True,
    fillable_mobile=True,
)


# Create and display a Shiny chat component
chat = ui.Chat(
    id="chat",
    messages=[
        "Hello! I am an assistant powered by Pydantic AI. Ask me anything you'd like to know or do.",
    ],
)
chat.ui()


# Generate a response when the user submits a message
@chat.on_user_submit
async def handle_user_input(user_input: str):
    stream = pydantic_stream_generator(user_input)
    await chat.append_message_stream(stream)


# An async generator function to stream the response from the Pydantic AI agent
async def pydantic_stream_generator(user_input: str):
    async with chat_client.run_stream(user_input) as result:
        async for chunk in result.stream_text(delta=True):
            yield chunk

from typing import List

from dotenv import load_dotenv
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage
from shiny import reactive
from shiny.express import ui

_ = load_dotenv()
chat_client = Agent(
    "openai:gpt-4.1-nano-2025-04-14",
    system_prompt="You are a helpful assistant.",
)

conversation_history = reactive.value(list[ModelMessage]([]))

ui.page_opts(
    title="Hello OpenAI Chat",
    fillable=True,
    fillable_mobile=True,
)

chat = ui.Chat(
    id="chat",
    messages=["Hello! How can I help you today?"],
)
chat.ui()


@chat.on_user_submit
async def handle_user_input(user_input: str):
    current_history = conversation_history.get()
    stream = pydantic_stream_generator(user_input, current_history)
    await chat.append_message_stream(stream)


async def pydantic_stream_generator(
    user_input: str, current_history: List[ModelMessage]
):
    message_history = current_history if current_history else None
    async with chat_client.run_stream(
        user_input, message_history=message_history
    ) as result:
        async for chunk in result.stream_text(delta=True):
            yield chunk
        conversation_history.set(result.all_messages())